{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 设置实验参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "mode = 'train'\n",
    "# mode = 'test'\n",
    "\n",
    "# folder = \"D:\\\\Downloads\\\\DATA\\\\水表\\\\标注\\\\8p\\\\8P blur\"\n",
    "folder = \"C:\\\\Users\\\\86198\\\\Desktop\\\\8针水表\\\\8针水表\\\\8P Uneven\"\n",
    "# folder = \"D:\\\\DATA\\\\multi-pointer-meter\\\\pointer-8\\\\8p uneven\"\n",
    "# folder = \"D:\\\\DATA\\\\multi-pointer-meter\\\\pointer-3\\\\3P uneven\"\n",
    "\n",
    "down_scale_factor = 8\n",
    "train_batch_sz = 2\n",
    "test_batch_sz = 1\n",
    "\n",
    "num_epochs = 10\n",
    "lr=0.01"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 载入库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'coco_eval'",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mModuleNotFoundError\u001B[0m                       Traceback (most recent call last)",
      "Cell \u001B[1;32mIn [4], line 27\u001B[0m\n\u001B[0;32m     25\u001B[0m \u001B[38;5;66;03m# sys.path.insert(0, './torchvision_det_references') #确保可以通过下面的语句导入位于子目录中的包\u001B[39;00m\n\u001B[0;32m     26\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mutils\u001B[39;00m\n\u001B[1;32m---> 27\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mtorchvision_det_references\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mengine\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m train_one_epoch, evaluate\n\u001B[0;32m     28\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtransforms\u001B[39;00m \u001B[38;5;28;01mas\u001B[39;00m \u001B[38;5;21;01mT\u001B[39;00m\n",
      "File \u001B[1;32m~\\Desktop\\pointer_meter\\torchvision_det_references\\engine.py:8\u001B[0m\n\u001B[0;32m      6\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorchvision\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mmodels\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mdetection\u001B[39;00m\u001B[38;5;21;01m.\u001B[39;00m\u001B[38;5;21;01mmask_rcnn\u001B[39;00m\n\u001B[0;32m      7\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mutils\u001B[39;00m\n\u001B[1;32m----> 8\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mcoco_eval\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m CocoEvaluator\n\u001B[0;32m      9\u001B[0m \u001B[38;5;28;01mfrom\u001B[39;00m \u001B[38;5;21;01mcoco_utils\u001B[39;00m \u001B[38;5;28;01mimport\u001B[39;00m get_coco_api_from_dataset\n\u001B[0;32m     12\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21mtrain_one_epoch\u001B[39m(model, optimizer, data_loader, device, epoch, print_freq, scaler\u001B[38;5;241m=\u001B[39m\u001B[38;5;28;01mNone\u001B[39;00m):\n",
      "\u001B[1;31mModuleNotFoundError\u001B[0m: No module named 'coco_eval'"
     ]
    }
   ],
   "source": [
    "import mlflow\n",
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision\n",
    "\n",
    "from torchvision.models.detection.faster_rcnn import FastRCNNPredictor\n",
    "from torchvision.models.detection import KeypointRCNN\n",
    "from torchvision.models.detection.rpn import AnchorGenerator\n",
    "from torchvision.models.detection.keypoint_rcnn import KeypointRCNNPredictor\n",
    "\n",
    "\n",
    "import torchvision.transforms as tT\n",
    "import torchvision.transforms.functional as F\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from PIL import Image, ExifTags\n",
    "import json\n",
    "import PySimpleGUI as sg\n",
    "from _pointer_meter_helpers import rotate_im_accord_exiftag, load_anno, load_valid_imfile_names, MeterPtDirDataset\n",
    "\n",
    "\n",
    "sys.path.insert(0, './torchvision_det_references') #确保可以通过下面的语句导入位于子目录中的包\n",
    "import utils\n",
    "from engine import train_one_epoch, evaluate\n",
    "import transforms as T\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 定义模型结构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 2 #our dataset has two classes only - background and meter\n",
    "num_keypoints = 2\n",
    "\n",
    "# load a pre-trained model for classification and return\n",
    "# only the features\n",
    "backbone = torchvision.models.mobilenet_v2(weights=\"DEFAULT\").features\n",
    "# FasterRCNN needs to know the number of\n",
    "# output channels in a backbone. For mobilenet_v2, it's 1280\n",
    "# so we need to add it here\n",
    "backbone.out_channels = 1280\n",
    "\n",
    "# let's make the RPN generate 5 x 3 anchors per spatial\n",
    "# location, with 5 different sizes and 3 different aspect\n",
    "# ratios. We have a Tuple[Tuple[int]] because each feature\n",
    "# map could potentially have different sizes and\n",
    "# aspect ratios\n",
    "anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),\n",
    "                                   aspect_ratios=((0.5, 1.0, 2.0),))\n",
    "\n",
    "# let's define what are the feature maps that we will\n",
    "# use to perform the region of interest cropping, as well as\n",
    "# the size of the crop after rescaling.\n",
    "# if your backbone returns a Tensor, featmap_names is expected to\n",
    "# be [0]. More generally, the backbone should return an\n",
    "# OrderedDict[Tensor], and in featmap_names you can choose which\n",
    "# feature maps to use.\n",
    "roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],\n",
    "                                                output_size=7,\n",
    "                                                sampling_ratio=2)\n",
    "\n",
    "keypoint_roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],\n",
    "                                                             output_size=14,\n",
    "                                                              sampling_ratio=2)\n",
    "\n",
    "# put the pieces together inside a FasterRCNN model\n",
    "model = KeypointRCNN(backbone,\n",
    "                   num_classes=num_classes,\n",
    "                   rpn_anchor_generator=anchor_generator,\n",
    "                   box_roi_pool=roi_pooler,\n",
    "                   keypoint_roi_pool=keypoint_roi_pooler)\n",
    "\n",
    "in_features = model.roi_heads.box_predictor.cls_score.in_features\n",
    "model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "in_features2 = model.roi_heads.keypoint_predictor.kps_score_lowres.in_channels\n",
    "model.roi_heads.keypoint_predictor = KeypointRCNNPredictor(in_features2, num_keypoints)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 准备数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_transform(train):\n",
    "    transforms = []\n",
    "    transforms.append(T.PILToTensor())\n",
    "    transforms.append(T.ConvertImageDtype(torch.float))\n",
    "    if train:\n",
    "        transforms.append(T.RandomPhotometricDistort())\n",
    "    return T.Compose(transforms)\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "dataset = MeterPtDirDataset(folder, get_transform(train=True), down_scale_factor)\n",
    "dataset_test = MeterPtDirDataset(folder, get_transform(train=False),down_scale_factor)\n",
    "\n",
    "# split the dataset in train and test set\n",
    "indices = torch.randperm(len(dataset)).tolist()\n",
    "dataset = torch.utils.data.Subset(dataset, indices[:-30])\n",
    "dataset_test = torch.utils.data.Subset(dataset_test, indices[-30:])\n",
    "\n",
    "# define training and validation data loaders\n",
    "data_loader = torch.utils.data.DataLoader(\n",
    "    dataset, batch_size=train_batch_sz, shuffle=True, num_workers=0,\n",
    "    collate_fn=utils.collate_fn)\n",
    "\n",
    "data_loader_test = torch.utils.data.DataLoader(\n",
    "    dataset_test, batch_size=test_batch_sz, shuffle=False, num_workers=0,\n",
    "    collate_fn=utils.collate_fn)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# construct an optimizer\n",
    "params = [p for p in model.parameters() if p.requires_grad]\n",
    "optimizer = torch.optim.SGD(params, lr,\n",
    "                            momentum=0.9, weight_decay=0.0005)\n",
    "# and a learning rate scheduler\n",
    "lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,\n",
    "                                                step_size=3,\n",
    "                                                gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'model' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn [3], line 2\u001B[0m\n\u001B[0;32m      1\u001B[0m fn \u001B[38;5;241m=\u001B[39m \u001B[38;5;124m'\u001B[39m\u001B[38;5;124mmodel_weights_xxx.pth\u001B[39m\u001B[38;5;124m'\u001B[39m\n\u001B[1;32m----> 2\u001B[0m torch\u001B[38;5;241m.\u001B[39msave(\u001B[43mmodel\u001B[49m\u001B[38;5;241m.\u001B[39mstate_dict(), fn) \n\u001B[0;32m      4\u001B[0m \u001B[38;5;28;01mif\u001B[39;00m mode \u001B[38;5;241m==\u001B[39m \u001B[38;5;124m\"\u001B[39m\u001B[38;5;124mtrain\u001B[39m\u001B[38;5;124m\"\u001B[39m:    \n\u001B[0;32m      5\u001B[0m     \u001B[38;5;66;03m# move model to the right device\u001B[39;00m\n\u001B[0;32m      6\u001B[0m     model\u001B[38;5;241m.\u001B[39mto(device)\n",
      "\u001B[1;31mNameError\u001B[0m: name 'model' is not defined"
     ]
    }
   ],
   "source": [
    "fn = 'model_weights_xxx.pth'\n",
    "torch.save(model.state_dict(), fn) \n",
    "\n",
    "if mode == \"train\":    \n",
    "    # move model to the right device\n",
    "    model.to(device)\n",
    "\n",
    "    # scaler = torch.cuda.amp.GradScaler()\n",
    "  \n",
    "    print_freq = 10\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        # train for one epoch, printing every 10 iterations\n",
    "        # train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler)\n",
    "        train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq)\n",
    "\n",
    "        # update the learning rate\n",
    "        lr_scheduler.step()\n",
    "        # evaluate on the test dataset\n",
    "        evaluate(model, data_loader_test, device=device)\n",
    "\n",
    "    print(\"Done!\")\n",
    "\n",
    "    #  save the model    \n",
    "    torch.save(model.state_dict(), fn) \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 分析为什么基于mobilenet v2的模型会那么大"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_param_num(Model):\n",
    "    total=sum([param.nelement() for param in Model.parameters()])\n",
    "    print(\"total param: {:.2f}M\".format(total/1e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total=sum([param.nelement() for param in model.parameters()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"total param: {:.2f}M\".format(total/1e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_backbone=sum([param.nelement() for param in backbone.parameters()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"total param: {:.2f}M\".format(total_backbone/1e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model.roi_heads.box_predictor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model.roi_heads.keypoint_predictor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(keypoint_roi_pooler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(roi_pooler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(anchor_generator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model.backbone)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model.rpn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model.roi_heads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_classes = 2 #our dataset has two classes only - background and meter\n",
    "num_keypoints = 2\n",
    "# load an instance segmentation model pre-trained on COCO\n",
    "model2 = torchvision.models.detection.keypointrcnn_resnet50_fpn(weights=\"DEFAULT\")\n",
    "\n",
    "in_features = model2.roi_heads.box_predictor.cls_score.in_features\n",
    "model2.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\n",
    "in_features2 = model2.roi_heads.keypoint_predictor.kps_score_lowres.in_channels\n",
    "model2.roi_heads.keypoint_predictor = KeypointRCNNPredictor(in_features2, num_keypoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model2.backbone)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model2.roi_heads)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print_param_num(model2.rpn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model2.backbone.out_channels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 单独调查resnet50和mobilenetv2的结构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "backbone_mobile = torchvision.models.mobilenet_v2(weights=\"DEFAULT\")\n",
    "backbone_resnet = torchvision.models.resnet50(weights=\"DEFAULT\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, m in enumerate(backbone_mobile.features.named_modules()):\n",
    "        print(idx, '->', m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, m in enumerate(backbone_resnet.named_modules()):\n",
    "        print(idx, '->', m)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "name": "myiste",
   "language": "python",
   "display_name": "my_test_jupyter"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "11938c6bc6919ae2720b4d5011047913343b08a43b18698fd82dedb0d4417594"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}